package com.jerome.election.storage;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Expiring map
 * Supports serialization
 *
 * @param <K> the type of keys maintained by this map
 * @param <V> the type of mapped values
 * @author Jerome
 * @date 2023/8/8
 */
public class ExpiringMap<K, V> implements Serializable {
    /**
     * Map to record key expiration times
     */
    private final Map<K, Long> expirationMap = new HashMap<>();
    /**
     * Map to store actual key-value pairs
     */
    private final Map<K, V> dataMap = new HashMap<>();
    /**
     * Default expiration time in milliseconds
     */
    private final long defaultExpiration;

    public ExpiringMap(long defaultExpiration) {
        this.defaultExpiration = defaultExpiration;
    }

    public synchronized void put(K key, V value) {
        put(key, value, defaultExpiration);
    }

    public synchronized void put(K key, V value, long expiration) {
        expirationMap.put(key, System.currentTimeMillis() + expiration);
        dataMap.put(key, value);
    }

    public synchronized V get(K key) {
        if (containsKey(key)) {
            if (isExpired(key)) {
                remove(key);
                return null;
            } else {
                return dataMap.get(key);
            }
        } else {
            return null;
        }
    }

    public synchronized boolean containsKey(K key) {
        return dataMap.containsKey(key);
    }

    /**
     * Remove expired keys
     *
     * @param key
     */
    public synchronized void removeExpired(K key) {
        if (isExpired(key)) {
            expirationMap.remove(key);
            dataMap.remove(key);
        }
    }

    public synchronized void remove(K key) {
        expirationMap.remove(key);
        dataMap.remove(key);
    }

    /**
     * Check if a key has expired
     *
     * @param key
     * @return true if the key has expired, false otherwise
     */
    private boolean isExpired(K key) {
        Long expirationTime = expirationMap.get(key);
        if (expirationTime != null) {
            return System.currentTimeMillis() > expirationTime;
        } else {
            return true;
        }
    }

    public synchronized void clear() {
        expirationMap.clear();
        dataMap.clear();
    }

    /**
     * Returns a {@link Set} view of the keys contained in this map.
     * The set is backed by the map, so changes to the map are
     * reflected in the set, and vice-versa.  If the map is modified
     * while an iteration over the set is in progress (except through
     * the iterator's own <tt>remove</tt> operation), the results of
     * the iteration are undefined.  The set supports element removal,
     * which removes the corresponding mapping from the map, via the
     * <tt>Iterator.remove</tt>, <tt>Set.remove</tt>,
     * <tt>removeAll</tt>, <tt>retainAll</tt>, and <tt>clear</tt>
     * operations.  It does not support the <tt>add</tt> or <tt>addAll</tt>
     * operations.
     *
     * @return a set view of the keys contained in this map
     */
    public Set<K> keySet(){
        return dataMap.keySet();
    }

    /**
     * Refresh the map and remove expired data
     *
     */
    public synchronized void refresh() {
        List<K> list = new ArrayList<>(dataMap.keySet());
        for (K key : list) {
            if (containsKey(key)) {
                if (isExpired(key)) {
                    remove(key);
                }
            }
        }
    }
}
